package com.inspur.edp.ai.flow.core.engine.databus;

import com.inspur.edp.ai.flow.engine.DataBus;
import com.inspur.edp.ai.flow.engine.StatefulSupporter;
import com.inspur.edp.ai.flow.exception.FlowException;
import io.iec.edp.caf.commons.utils.StringUtils;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class DataBusManager implements StatefulSupporter {

    private final Map<String, DataBusHolder> dataBusMap;
    // 锁用于保证某些操作的原子性
    private final Lock cleanupLock = new ReentrantLock();

    // 可选的清理机制相关字段
    private final long defaultTtl; // 默认存活时间(毫秒)，-1表示永不过期

    private final boolean enableCleanup; // 是否启用自动清理

    /**
     * 默认配置
     */
    public DataBusManager() {
        this(60000, true);
    }

    /**
     * 构造函数
     *
     * @param defaultTtl    @param defaultTtl 默认存活时间(毫秒)，-1表示永不过期
     * @param enableCleanup enableCleanup 是否启用自动清理过期上下文
     */
    public DataBusManager(long defaultTtl, boolean enableCleanup) {
        this.dataBusMap = new ConcurrentHashMap<>();
        this.defaultTtl = defaultTtl;
        this.enableCleanup = enableCleanup;
    }

    /**
     * 备份上下文
     *
     * @param instanceId 实例ID
     * @param dataBus    数据总线
     * @
     */
    @Override
    public void backup(String instanceId, DataBus dataBus) {
        backup(instanceId, dataBus, defaultTtl);
    }

    /**
     * 备份上下文（自定义TTL）
     *
     * @param instanceId 实例ID
     * @param dataBus    上下文对象
     * @param ttl        存活时间(毫秒)
     */
    public void backup(String instanceId, DataBus dataBus, long ttl) {
        if (StringUtils.isEmpty(instanceId)) {
            throw new FlowException("instanceId cannot be null or empty");
        }

        if (dataBus == null) {
            throw new FlowException("dataBus cannot be null");
        }

        DataBusHolder holder = new DataBusHolder(dataBus, ttl);
        dataBusMap.put(instanceId, holder);

        if (enableCleanup) {
            cleanupExpired();
        }
    }

    /**
     * 恢复DataBus
     *
     * @param instanceId 实例ID
     * @return 上下文对象，如果不存在或已过期则返回null
     */
    @Override
    public DataBus recovery(String instanceId) {
        if (instanceId == null || instanceId.trim().isEmpty()) {
            throw new FlowException("instanceId cannot be null or empty");
        }

        DataBusHolder holder = dataBusMap.get(instanceId);

        if (holder == null) {
            return null; // 不存在
        }

        // 检查是否过期
        if (holder.isExpired()) {
            dataBusMap.remove(instanceId); // 移除过期项目
            return null;
        }

        return holder.getDataBus();
    }

    /**
     * 移除指定上下文
     *
     * @param instanceId 实例ID
     * @return 被移除的上下文，如果不存在则返回null
     */
    public DataBus remove(String instanceId) {
        DataBusHolder holder = dataBusMap.remove(instanceId);
        return holder != null ? holder.getDataBus() : null;
    }

    /**
     * 检查是否存在指定实例ID的上下文
     *
     * @param instanceId 实例ID
     * @return 是否存在且未过期
     */
    public boolean contains(String instanceId) {
        DataBusHolder holder = dataBusMap.get(instanceId);
        return holder != null && !holder.isExpired();
    }

    /**
     * 清理所有过期的DataBus
     */
    public void cleanupExpired() {
        // 使用锁避免多个线程同时执行清理（非必须，但可以减少重复清理的开销）
        if (cleanupLock.tryLock()) {
            try {
                dataBusMap.entrySet().removeIf(entry -> {
                    DataBusHolder holder = entry.getValue();
                    return holder.isExpired();
                });
            } finally {
                cleanupLock.unlock();
            }
        }
    }

    /**
     * 获取当前存储的上下文数量（包括已过期的）
     *
     * @return 上下文数量
     */
    public int size() {
        return dataBusMap.size();
    }

    /**
     * 获取未过期的上下文数量
     *
     * @return 未过期的上下文数量
     */
    public int validSize() {
        return (int) dataBusMap.values().stream()
                .filter(holder -> !holder.isExpired())
                .count();
    }

    /**
     * 清空所有上下文
     */
    public void clear() {
        dataBusMap.clear();
    }
}
